import numpy as np
from scipy import sparse

# coefficient parameters for reaction and diffusion equation
def func_a(x, y):
    return 1 + 2 * y**2

def func_r(x, y):
    return 1 + x ** 2

# Assemble linear system for Laplace equation
def lap_A(n):
    N2 = n**2
    A = sparse.lil_matrix((N2, N2))
    for i in range(1, n-1):
        for j in range(1, n-1):
            idx = i * n + j
            A[idx, idx] += 4
            A[idx, idx-1] = -1
            A[idx, idx+1] = -1
            A[idx, idx-n] = -1
            A[idx, idx+n] = -1
            
    # Homogeneous Dirichlet Boundary
    for i in range(0, n):
        idx = 0 * n + i
        A[idx, idx] = 1

        idx = (n-1) * n + i
        A[idx, idx] = 1

        idx = i * n
        A[idx, idx] = 1

        idx = i * n + n - 1
        A[idx, idx] = 1
    A = A.tocsc()
    return A

def lap_b(f, h):
    h2 = h * h
    b = np.zeros_like(f)
    b[1:-1, 1:-1] = f[1:-1, 1:-1]
    b = b * h2
    return b.reshape(-1)

# Assemble linear system for reaction equation
def reaction_A(n):
    n2 = n**2
    h = 2 / (n - 1)
    A = np.identity(n2)
    x = np.linspace(-1, 1, n)
    y = np.linspace(-1, 1, n)
    xx, yy = np.meshgrid(x, y)

    a = func_a(xx, yy)
    r = func_r(xx, yy)

    for i in range(1, n-1):
        for j in range(1, n-1):
            idx = i * n + j
            A[idx, idx]   = 4 * a[i, j] + a[i, j+1] + a[i, j-1] + a[i+1, j] + a[i-1, j] + 2 * h**2 * r[i, j]
            A[idx, idx-1] = -(a[i, j] + a[i, j-1]) 
            A[idx, idx+1] = -(a[i, j] + a[i, j+1])
            A[idx, idx-n] = -(a[i, j] + a[i-1, j]) 
            A[idx, idx+n] = -(a[i, j] + a[i+1, j])
    return sparse.csc_matrix(A)

def reaction_b(f, h):
    h2 = 2 * h**2
    # h2 = h**2
    b = np.zeros_like(f)
    b[1:-1, 1:-1] = f[1:-1, 1:-1]
    b = b * h2
    return b.reshape(-1)


def l2_norm(a, b, h):
    return np.sqrt(((a - b)**2 * h**2).sum())

def cal_l2(pres, U, n):
    h = 2 / (n - 1)
    l2_errors = []
    for i in range(len(U)):
        pre, ans = pres[i].squeeze(), U[i].squeeze()
        l2 = l2_norm(ans, pre, h)
        l2_errors.append(l2)
    return l2_errors